from BinaryTree import BinaryTree
from NumMatrix import NumMatrix

class HierarchicalClustering:

    def __init__(self, matdists:NumMatrix):
        self.matdists = matdists
    
    def execute_clustering(self) -> BinaryTree:
        '''
        weighted average of the distances:\n
        D(A or B,X) = { |A|*D(A,X)+|B|*D(B,X) } / { |A|+|B| }
        '''
        trees = []
        tableDist = self.matdists.copy()
        for i in range(self.matdists.num_rows()):
            t = BinaryTree(i)
            trees.append(t)
        
        while len(trees) > 1 :
            i,j = tableDist.min_dist_indexes()
            child = BinaryTree(-1, tableDist.get_value(i, j)/2.0, trees[i], trees[j])
            ti = trees.pop(i)
            tj = trees.pop(j)
            dists = self.dist_update(tableDist,ti,tj,(i,j))
            tableDist.remove_row(i)
            tableDist.remove_row(j)
            if not tableDist.is_empty():
                tableDist.remove_col(i)
                tableDist.remove_col(j)
            tableDist.add_row(dists)
            tableDist.add_col([0] * (len(dists)+1))
            trees.append(child)
        return trees[0]

    def dist_update(_,tableDist:NumMatrix,ti:BinaryTree,tj:BinaryTree,ij) -> list:
        i,j = ij
        dists = []
        for x in range(tableDist.num_rows()):          
            if x != i and x != j:
                si = len(ti.get_cluster())
                sj = len(tj.get_cluster())
                d = (si*tableDist.get_value(i,x) + sj*tableDist.get_value(j,x)) / (si+sj)
                dists.append(d)
        return dists


def test():
    m = NumMatrix(5,5)
    m.set_value(0, 1, 2)
    m.set_value(0, 2, 5)
    m.set_value(0, 3, 7)
    m.set_value(0, 4, 9)
    m.set_value(1, 2, 4)
    m.set_value(1, 3, 6)
    m.set_value(1, 4, 7)
    m.set_value(2, 3, 4)
    m.set_value(2, 4, 6)
    m.set_value(3, 4, 3)
    hc = HierarchicalClustering(m)
    arv = hc.execute_clustering()
    arv.print_tree()
    
if __name__ == '__main__': 
    test()
